import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm

# ---------------- 配置全局参数 ----------------
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC", "Times New Roman"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号
plt.rcParams.update({
    "font.size": 12,
    "axes.labelsize": 14,
    "axes.titlesize": 16,
    "legend.fontsize": 11,
    "grid.linestyle": "--",
    "grid.alpha": 0.7
})

# 尝试启用 LaTeX，如果不可用则忽略
try:
    plt.rcParams["text.usetex"] = True
    plt.rcParams["text.latex.preamble"] = [
        r"\usepackage{amsmath}",
        r"\usepackage{amsfonts}",
        r"\usepackage{amsbsy}"
    ]
except Exception:
    plt.rcParams["text.usetex"] = False


def plot_drone_trajectory_3d():
    """绘制无人机运动轨迹三维示意图"""

    # ---------------- 参数定义 ----------------
    x0, y0, z0 = 17800, 0, 1800  # 初始位置
    t_total, n_points = 30, 200  # 总时长与采样点
    t_values = np.linspace(0, t_total, n_points)
    annotate_times = [10, 20, 30]

    # 方位角θ（度）和仰角φ（度）
    thetas = [0, 90, 180, 270]
    phis = [-30, 0, 30]

    # 无人机速度（m/s）
    velocities = [70, 105, 140]

    # ---------------- 绘图 ----------------
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, projection="3d")

    # 使用渐变色区分速度
    colors = cm.viridis(np.linspace(0, 0.9, len(velocities)))

    for v_idx, v in enumerate(velocities):
        linewidth = 2.5 - v_idx * 0.5  # 不同速度线宽
        color = colors[v_idx]

        for theta in thetas:
            for phi in phis:
                theta_rad, phi_rad = np.radians(theta), np.radians(phi)

                # 计算轨迹
                x = x0 + v * t_values * np.cos(phi_rad) * np.cos(theta_rad)
                y = y0 + v * t_values * np.cos(phi_rad) * np.sin(theta_rad)
                z = z0 + v * t_values * np.sin(phi_rad)

                alpha = 0.9 if phi == 0 else 0.6
                label = f"$v_U = {v}$ m/s" if (theta == 0 and phi == 0) else None

                ax.plot(x, y, z, color=color, linewidth=linewidth,
                        alpha=alpha, label=label)

                # 标注特定时间点
                for t in annotate_times:
                    t_idx = np.argmin(np.abs(t_values - t))
                    marker_size = 60 if phi == 0 else 40

                    ax.scatter(x[t_idx], y[t_idx], z[t_idx],
                               color=color, s=marker_size,
                               edgecolors="black", linewidth=0.8, zorder=5)

                    if phi == 0 and (theta in [0, 180]):
                        ax.text3D(x[t_idx], y[t_idx], z[t_idx],
                                  f"$t={t}$s", color=color,
                                  fontsize=10, weight="bold")

        # 标记每条轨迹的终点
        ax.scatter(x[-1], y[-1], z[-1], color=color,
                   marker="^", s=100, edgecolors="black",
                   linewidth=1.0, zorder=6)

    # ---------------- 初始点标记 ----------------
    ax.scatter(x0, y0, z0, color="red", s=150, marker="o",
               edgecolors="black", linewidth=1.5, zorder=10,
               label="初始位置")
    ax.text3D(x0, y0, z0, f"$({x0}, {y0}, {z0})$ m",
              color="red", fontsize=12, weight="bold")

    # ---------------- 图形美化 ----------------
    ax.set_xlabel("$x$ 坐标 (m)", labelpad=15)
    ax.set_ylabel("$y$ 坐标 (m)", labelpad=15)
    ax.set_zlabel("$z$ 坐标 (m)", labelpad=15)
    ax.set_title("无人机运动轨迹三维示意图", fontsize=18, pad=20)

    ax.legend(loc="upper left", bbox_to_anchor=(1.02, 1),
              frameon=True, framealpha=0.9, edgecolor="gray")

    # 添加角度说明
    angle_text = (r"方位角 $\theta$: 0°, 90°, 180°, 270°" + "\n" +
                  r"仰角 $\phi$: -30°, 0°, 30°")
    plt.figtext(0.01, 0.01, angle_text, fontsize=12,
                bbox=dict(facecolor="white", alpha=0.8, boxstyle="round,pad=0.5"))

    # 添加图注
    plt.figtext(0.5, 0.005,
                "图注：以FY1初始位置(17800, 0, 1800)m为起点，"
                "展示不同方位角θ与仰角φ组合下的飞行轨迹，"
                "标注t=10s、20s、30s时的位置坐标，"
                "用颜色区分不同速度v_U的轨迹差异",
                ha="center", fontsize=11,
                bbox=dict(facecolor="white", alpha=0.8, boxstyle="round,pad=0.5"))

    # 视角 & 网格
    ax.view_init(elev=30, azim=45)
    ax.grid(True)

    # 布局与保存
    plt.tight_layout()
    plt.savefig("drone_trajectory_3d.png", dpi=300, bbox_inches="tight")
    plt.show()


if __name__ == "__main__":
    plot_drone_trajectory_3d()
